{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gGxCD4mGHHjG"
   },
   "source": [
    "# Ungraded Lab: Data Augmentation\n",
    "\n",
    "In the previous lessons, you saw that having a high training accuracy does not automatically mean having a good predictive model. It can still perform poorly on new data because it has overfit to the training set. In this lab, you will see how to avoid that using _data augmentation_. This increases the amount of training data by modifying the existing training data's properties. For example, in image data, you can apply different preprocessing techniques such as rotate, flip, shear, or zoom on your existing images so you can simulate other data that the model should also learn from. This way, the model would see more variety in the images during training so it will infer better on new, previously unseen data.\n",
    "\n",
    "Let's see how you can do this in the following sections."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kJJqX4DxcQs8"
   },
   "source": [
    "## Baseline Performance\n",
    "\n",
    "You will start with a model that's very effective at learning `Cats vs Dogs` without data augmentation. It's similar to the previous models that you have used. Note that there are four convolutional layers with 32, 64, 128 and 128 convolutions respectively. The code is basically the same from the previous lab so we won't go over the details step by step since you've already seen it before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First import the necessary libraries\n",
    "import os\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_DyUfCTgdwa8"
   },
   "outputs": [],
   "source": [
    "BASE_DIR = 'cats_and_dogs_filtered'\n",
    "\n",
    "train_dir = os.path.join(BASE_DIR, 'train')\n",
    "validation_dir = os.path.join(BASE_DIR, 'validation')\n",
    "\n",
    "# Directory with training cat/dog pictures\n",
    "train_cats_dir = os.path.join(train_dir, 'cats')\n",
    "train_dogs_dir = os.path.join(train_dir, 'dogs')\n",
    "\n",
    "# Directory with validation cat/dog pictures\n",
    "validation_cats_dir = os.path.join(validation_dir, 'cats')\n",
    "validation_dogs_dir = os.path.join(validation_dir, 'dogs')\n",
    "\n",
    "print(f\"Contents of base directory: {os.listdir(BASE_DIR)}\")\n",
    "\n",
    "print(f\"\\nContents of train directory: {os.listdir(train_dir)}\")\n",
    "\n",
    "print(f\"\\nContents of validation directory: {os.listdir(validation_dir)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ub_BdOJIfZ_Q"
   },
   "source": [
    "You will place the model creation inside a function so you can easily initialize a new one when you use data augmentation later in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uWllK_Wad-Mx"
   },
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    '''Creates a CNN with 4 convolutional layers'''\n",
    "    model = tf.keras.models.Sequential([\n",
    "        tf.keras.Input(shape=(150, 150, 3)),\n",
    "        tf.keras.layers.Rescaling(1./255),\n",
    "        tf.keras.layers.Conv2D(32, (3,3), activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(2, 2),\n",
    "        tf.keras.layers.Conv2D(64, (3,3), activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(2,2),\n",
    "        tf.keras.layers.Conv2D(128, (3,3), activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(2,2),\n",
    "        tf.keras.layers.Conv2D(128, (3,3), activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(2,2),\n",
    "        tf.keras.layers.Flatten(),\n",
    "        tf.keras.layers.Dense(512, activation='relu'),\n",
    "        tf.keras.layers.Dense(1, activation='sigmoid')\n",
    "    ])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Uqe1AyzAavoP"
   },
   "source": [
    "You will preprocess the training and validation datasets as usual."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3sj3ERldRGpY"
   },
   "outputs": [],
   "source": [
    "# Instantiate the training dataset\n",
    "train_dataset = tf.keras.utils.image_dataset_from_directory(\n",
    "    train_dir,\n",
    "    image_size=(150, 150),\n",
    "    batch_size=20,\n",
    "    label_mode='binary'\n",
    "    )\n",
    "\n",
    "# Instantiate the validation dataset\n",
    "validation_dataset = tf.keras.utils.image_dataset_from_directory(\n",
    "    validation_dir,\n",
    "    image_size=(150, 150),\n",
    "    batch_size=20,\n",
    "    label_mode='binary'\n",
    "    )\n",
    "\n",
    "# Optimize the datasets for training\n",
    "SHUFFLE_BUFFER_SIZE = 1000\n",
    "PREFETCH_BUFFER_SIZE = tf.data.AUTOTUNE\n",
    "\n",
    "train_dataset_final = (train_dataset\n",
    "                       .cache()\n",
    "                       .shuffle(SHUFFLE_BUFFER_SIZE)\n",
    "                       .prefetch(PREFETCH_BUFFER_SIZE)\n",
    "                       )\n",
    "\n",
    "validation_dataset_final = (validation_dataset\n",
    "                            .cache()\n",
    "                            .prefetch(PREFETCH_BUFFER_SIZE)\n",
    "                            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ig2aYb-ZbkY2"
   },
   "source": [
    "You will train only for 20 epochs to save time but feel free to increase this if you want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hdqUoF44esR3"
   },
   "outputs": [],
   "source": [
    "# Constant for epochs\n",
    "EPOCHS = 20\n",
    "\n",
    "# Create a new model\n",
    "model = create_model()\n",
    "\n",
    "# Setup the training parameters\n",
    "model.compile(loss='binary_crossentropy',\n",
    "              optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-4),\n",
    "              metrics=['accuracy'])\n",
    "\n",
    "# Train the model\n",
    "history = model.fit(\n",
    "      train_dataset_final,\n",
    "      epochs=EPOCHS,\n",
    "      validation_data=validation_dataset_final,\n",
    "      verbose=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y-G0Am4cguNt"
   },
   "source": [
    "You will then visualize the loss and accuracy with respect to the training and validation set. You will again use a convenience function so it can be reused later. This function accepts a [History](https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/History) object which contains the results of the `fit()` method you ran above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GZWPcmKWO303"
   },
   "outputs": [],
   "source": [
    "def plot_loss_acc(history):\n",
    "    '''Plots the training and validation loss and accuracy from a history object'''\n",
    "    acc = history.history['accuracy']\n",
    "    val_acc = history.history['val_accuracy']\n",
    "    loss = history.history['loss']\n",
    "    val_loss = history.history['val_loss']\n",
    "\n",
    "    epochs = range(len(acc))\n",
    "\n",
    "    fig, ax = plt.subplots(1,2, figsize=(12, 6))\n",
    "    ax[0].plot(epochs, acc, 'bo', label='Training accuracy')\n",
    "    ax[0].plot(epochs, val_acc, 'b', label='Validation accuracy')\n",
    "    ax[0].set_title('Training and validation accuracy')\n",
    "    ax[0].set_xlabel('epochs')\n",
    "    ax[0].set_ylabel('accuracy')\n",
    "    ax[0].legend()\n",
    "\n",
    "    ax[1].plot(epochs, loss, 'bo', label='Training Loss')\n",
    "    ax[1].plot(epochs, val_loss, 'b', label='Validation Loss')\n",
    "    ax[1].set_title('Training and validation loss')\n",
    "    ax[1].set_xlabel('epochs')\n",
    "    ax[1].set_ylabel('loss')\n",
    "    ax[1].legend()\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Vojz4NYXiT_f"
   },
   "outputs": [],
   "source": [
    "# Plot training results\n",
    "plot_loss_acc(history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zb81GvNov-Tg"
   },
   "source": [
    "From the results above, you'll see the training accuracy is more than 90%, and the validation accuracy is in the 70%-80% range. This is a great example of _overfitting_ -- which in short means that it can do very well with images it has seen before, but not so well with images it hasn't.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5KBz-vFbjLZX"
   },
   "source": [
    "## Data augmentation\n",
    "\n",
    "One simple method to avoid overfitting is to augment the images. If you think about it, most pictures of a cat are very similar -- the ears are at the top, the eyes are below the ears etc. Things like the distance between the eyes and ears will always be quite similar too.\n",
    "\n",
    "What if you tweak with the images a bit -- rotate the image, squash it, etc.  That's what image augementation is all about.\n",
    "\n",
    "To do that, you will build a data augmentation model with [preprocessing layers for image augmentation](https://www.tensorflow.org/guide/keras/preprocessing_layers#image_data_augmentation). This will transform the data during training to introduce variations of the same image. Let's quickly go over the layers you will use in this exercise.\n",
    "\n",
    "* [RandomFlip](https://www.tensorflow.org/api_docs/python/tf/keras/layers/RandomFlip) is for randomly flipping the images horizontally, vertically, or both.\n",
    "* [RandomRotation](https://www.tensorflow.org/api_docs/python/tf/keras/layers/RandomRotation) rotates the image by an angle within a given range.\n",
    "* [RandomTranslation](https://www.tensorflow.org/api_docs/python/tf/keras/layers/RandomTranslation) shifts pictures vertically and horizontally.\n",
    "* [RandomZoom](https://www.tensorflow.org/api_docs/python/tf/keras/layers/RandomZoom) zooms into or out of the images.\n",
    "\n",
    "In addition, some of these layers have a `fill_mode` parameter. This is the strategy used for filling newly created pixels, which can appear after a rotation or a width/height shift.\n",
    "\n",
    "The code below will create this model with some set parameters. After you complete this lab, feel free to modify these and see the impact on the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FUv0b0TCRzJN"
   },
   "outputs": [],
   "source": [
    "# Define fill mode.\n",
    "FILL_MODE = 'nearest'\n",
    "\n",
    "# Create the augmentation model.\n",
    "data_augmentation = tf.keras.Sequential([\n",
    "    # Specify the input shape.\n",
    "    tf.keras.Input(shape=(150,150,3)),\n",
    "    # Add the augmentation layers\n",
    "    tf.keras.layers.RandomFlip(\"horizontal\"),\n",
    "    tf.keras.layers.RandomRotation(0.2, fill_mode=FILL_MODE),\n",
    "    tf.keras.layers.RandomTranslation(0.2,0.2, fill_mode=FILL_MODE),\n",
    "    tf.keras.layers.RandomZoom(0.2, fill_mode=FILL_MODE)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T5xj5HYXg6mJ"
   },
   "source": [
    "You will define a utility function that lets you preview how the transformed images look like. It will take in a sample image, then output a given number of augmented images using the model you defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z0Q9C8hP4P6O"
   },
   "outputs": [],
   "source": [
    "def demo_augmentation(sample_image, model, num_aug):\n",
    "    '''Takes a single image array, then uses a model to generate num_aug transformations'''\n",
    "\n",
    "    # Instantiate preview list\n",
    "    image_preview = []\n",
    "\n",
    "    # Convert input image to a PIL image instance\n",
    "    sample_image_pil = tf.keras.utils.array_to_img(sample_image)\n",
    "\n",
    "    # Append the result to the list\n",
    "    image_preview.append(sample_image_pil)\n",
    "\n",
    "    # Apply the image augmentation and append the results to the list\n",
    "    for i in range(NUM_AUG):\n",
    "        sample_image_aug = model(tf.expand_dims(sample_image, axis=0))\n",
    "        sample_image_aug_pil = tf.keras.utils.array_to_img(tf.squeeze(sample_image_aug))\n",
    "        image_preview.append(sample_image_aug_pil)\n",
    "\n",
    "    # Instantiate a subplot\n",
    "    fig, axes = plt.subplots(1, NUM_AUG + 1, figsize=(12, 12))\n",
    "\n",
    "    # Preview the images.\n",
    "    for index, ax in enumerate(axes):\n",
    "        ax.imshow(image_preview[index])\n",
    "        ax.set_axis_off()\n",
    "\n",
    "        if index == 0:\n",
    "            ax.set_title('original')\n",
    "        else:\n",
    "            ax.set_title(f'augment {index}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O-9z1jW_ijo2"
   },
   "source": [
    "Now get some images from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ap0OnQKfX0MI"
   },
   "outputs": [],
   "source": [
    "# Get a batch of images\n",
    "sample_batch = list(train_dataset.take(1))[0][0]\n",
    "print(f'images per batch: {len(sample_batch)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OU9R7e9FirkZ"
   },
   "source": [
    "This will show sample transformations for the first 4 images of the sample batch. Notice that each call of the `data_augmentation` model yields a different output. It's like adding more images to your dataset without you having to collect them manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iBIbpv7GXX-h"
   },
   "outputs": [],
   "source": [
    "NUM_AUG = 4\n",
    "\n",
    "# Apply the transformations to the first 4 images\n",
    "demo_augmentation(sample_batch[0], data_augmentation, NUM_AUG)\n",
    "demo_augmentation(sample_batch[1], data_augmentation, NUM_AUG)\n",
    "demo_augmentation(sample_batch[2], data_augmentation, NUM_AUG)\n",
    "demo_augmentation(sample_batch[3], data_augmentation, NUM_AUG)\n",
    "\n",
    "# Uncomment the line below to delete the variable to free up some memory\n",
    "# del sample_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B9rZ1bTIjjIP"
   },
   "source": [
    "Now that you see what the preprocessing layers do, you can prepend these to the base model so it can generate transformed images to the base model. Do note that these layers are only active while training. They are automatically disabled during prediction and evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OXVW2NzMkW_i"
   },
   "outputs": [],
   "source": [
    "# Instantiate the base model\n",
    "model_without_aug = create_model()\n",
    "\n",
    "# Prepend the data augmentation layers to the base model\n",
    "model_with_aug = tf.keras.models.Sequential([\n",
    "    data_augmentation,\n",
    "    model_without_aug\n",
    "])\n",
    "\n",
    "# Compile the model\n",
    "model_with_aug.compile(\n",
    "    loss='binary_crossentropy',\n",
    "    optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-4),\n",
    "    metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "maFVjXG9lFY8"
   },
   "source": [
    "Because you now have virtually more data, it will also take the model more time to learn the relevant features. Without data augmentation, your model already started overfitting to the training set within 20 epochs. Try training this model for 80 epochs and observe the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UK7_Fflgv8YC"
   },
   "outputs": [],
   "source": [
    "EPOCHS=80\n",
    "\n",
    "# Train the new model\n",
    "history_with_aug = model_with_aug.fit(\n",
    "      train_dataset_final,\n",
    "      epochs=EPOCHS,\n",
    "      validation_data=validation_dataset_final,\n",
    "      verbose=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bnyRnwopT5aW"
   },
   "outputs": [],
   "source": [
    "# Plot the results of training with data augmentation\n",
    "plot_loss_acc(history_with_aug)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1D1hd5fqmJUx"
   },
   "source": [
    "As you can see, the training accuracy has gone down compared to the baseline. This is expected because (as a result of data augmentation) there are more variety in the images so the model will need more runs to learn from them. The good thing is the validation accuracy is no longer stalling and is more in line with the training results. This means that the model is now performing better on unseen data.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z4B9b6GPnKg1"
   },
   "source": [
    "## Wrap Up\n",
    "\n",
    "This exercise showed a simple trick to avoid overfitting. You can improve your baseline results by simply tweaking the same images you have already. The image augmentation preprocessing layers will do just that. Try to modify the values in the augmentation model and see what results you get. You can also add other preprocessing layers for random contrast, brightness, or cropping.\n",
    "\n",
    "Take note that this will not work for all cases. In the next lesson, Laurence will show a scenario where data augmentation will not help improve your validation accuracy.\n",
    "\n",
    "Before going back to the classroom, run the cell below to free up resources for the next lab. You might see a pop-up about restarting the kernel afterwards. You can safely ignore it and just press Ok. You can then close this lab, then go back to the classroom for the next lecture. See you there!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Shutdown the kernel to free up resources. \n",
    "# Note: You can expect a pop-up when you run this cell. You can safely ignore that and just press `Ok`.\n",
    "\n",
    "from IPython import get_ipython\n",
    "\n",
    "k = get_ipython().kernel\n",
    "\n",
    "k.do_shutdown(restart=False)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
